package org.dromara.easyai.extensions.ascend;

import org.dromara.easyai.extensions.ascend.core.JCannxException;
import org.dromara.easyai.matrixTools.AscendMatrix;
import org.dromara.easyai.matrixTools.Matrix;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * @author: wen.y
 * @date: 2025/6/19
 */
public class AscendMatrixImpl implements AscendMatrix {
	private static final Logger logger = LoggerFactory.getLogger(AscendMatrixImpl.class);

	static {
		logger.info("EasyAI Ascend extensions loaded.");
	}

	@Override
	public void init() throws Exception {
		try {
			JCannx.initialize();
			int initRet = JCannx.init();
			logger.info("EasyAI Ascend extensions JCann init ret={}", initRet);
			if (initRet != 0) {
				throw new JCannxException("EasyAI Ascend extensions init failed.");
			}
			logger.info("EasyAI Ascend extensions init success.");
		} catch (Throwable e) {
			logger.error("EasyAI Ascend extensions init failed.", e);
			throw new JCannxException("EasyAI Ascend extensions init error.");
		}

	}

	@Override
	public Matrix mulMatrix(Matrix a, Matrix b) throws Exception {
		float[] h_A = a.getCudaMatrix();
		float[] h_B = b.getCudaMatrix();
		float[] h_C = new float[a.getX() * b.getY()];

//		System.out.println("Java Input h_A:");
//		for (float v : h_A) {
//			System.out.print(v + " ");
//		}
//		System.out.println("\nJava Input h_B:");
//		for (float v : h_B) {
//			System.out.print(v + " ");
//		}
//		System.out.println();

		int ret = JCannx.matrixMul(h_A, h_B, h_C, a.getX(), a.getY(), b.getY());

//		System.out.println("ret:" + ret);
//		// 打印结果数据 (Java端)
//		System.out.println("Java Result h_C:");
//		for (float v : h_C) {
//			System.out.print(v + " ");
//		}
//		System.out.println();

		if (ret != 0) {
			throw new JCannxException("matrixMul error " + ret);
		}

		Matrix res = new Matrix(a.getX(), b.getY());
		res.setCudaMatrix(h_C, a.getX(), b.getY());
		return res;
	}
}
